# 导入必要的库
import torch
import cv2
import numpy as np
import os
import glob as glob

# 从xml.etree导入ElementTree用于解析XML文件
from xml.etree import ElementTree as et
# 从config导入配置项
# 导入PyTorch的数据集和数据加载器
from torch.utils.data import Dataset, DataLoader
# 导入自定义工具函数


# 定义数据集类
class CustomDataset(Dataset):
    def __init__(self, image_paths, mask_paths,width, height, classes, transforms=None):
        # 初始化数据集
        self.transforms = transforms  # 图像变换
        self.height = height  # 图像调整后的高度
        self.width = width  # 图像调整后的宽度
        self.classes = classes  # 类别列表

        # 获取所有图像路径，并按顺序排列
        self.image_paths = image_paths
        self.mask_paths = mask_paths

    def __getitem__(self, idx):
        # 根据索引获取图像
        image_path = self.image_paths[idx]
        annot_file_path = self.mask_paths[idx]

        # 读取图像
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)  # BGR转RGB
        image_resized = cv2.resize(image, (self.width, self.height))  # 调整图像大小
        image_resized /= 255.0  # 归一化

        boxes = []  # 边界框
        labels = []  # 标签
        tree = et.parse(annot_file_path)  # 解析XML
        root = tree.getroot()

        # 提取图像原始尺寸
        image_width = image.shape[1]
        image_height = image.shape[0]

        # 提取并调整边界框尺寸
        for member in root.findall('object'):
            labels.append(self.classes.index(member.find('name').text))  # 获取标签索引

            xmin = int(member.find('bndbox').find('xmin').text)
            xmax = int(member.find('bndbox').find('xmax').text)
            ymin = int(member.find('bndbox').find('ymin').text)
            ymax = int(member.find('bndbox').find('ymax').text)

            xmin_final = (xmin / image_width) * self.width
            xmax_final = (xmax / image_width) * self.width
            ymin_final = (ymin / image_height) * self.height
            ymax_final = (ymax / image_height) * self.height

            boxes.append([xmin_final, ymin_final, xmax_final, ymax_final])

        # 将边界框和标签转换为张量
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])  # 边界框面积
        iscrowd = torch.zeros((boxes.shape[0],), dtype=torch.int64)  # 无众包实例
        labels = torch.as_tensor(labels, dtype=torch.int64)

        # 准备最终的`target`字典
        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["area"] = area
        target["iscrowd"] = iscrowd
        image_id = torch.tensor([idx])
        target["image_id"] = image_id

        # 应用图像变换
        if self.transforms:
            sample = self.transforms(image=image_resized,
                                     bboxes=target['boxes'],
                                     labels=labels)
            image_resized = sample['image']
            target['boxes'] = torch.Tensor(sample['bboxes'])

        return image_resized, target

    def __len__(self):
        # 返回数据集中图像的总数
        return len(self.image_paths)






